# import torch,torchvision
import torch
from torch import nn

# from torch.utils.data import DataLoader
#
device = torch.device("cuda")
#
class my_Vggnet(nn.Module):
    def __init__(self,num_classees=1000):
        super(my_Vggnet, self).__init__()
        layer= []
        in_dim=3
        out_dim=64
        for i in range(13):
            layer+=[nn.Conv2d(in_channels=in_dim,out_channels=out_dim,kernel_size=(3,3),stride=(1,1),padding=1),nn.ReLU(inplace=True)]
            in_dim=out_dim
            if i==1 or i==3 or i==6 or i==9 or i==12:
                layer+=[nn.MaxPool2d(2,2)]
                if i!=9:
                    out_dim*=2
        self.feature=nn.Sequential(*layer)
        self.classifier=nn.Sequential(
            nn.Linear(512,4096),
            nn.ReLU(True),
            nn.Dropout(),
            nn.Linear(4096,4096),
            nn.ReLU(True),
            nn.Dropout(),
            nn.Linear(4096,num_classees),
        )
    def forward(self,x):
        x=self.feature(x)
        # print(x.shape)
        x = x.view(x.size(0), -1)
        # print(x.shape)
        x=self.classifier(x)
        return x
#
my_module=my_Vggnet()
# my_module.to(device)
# # loss_fn = nn.CrossEntropyLoss()
# # loss_fn.to(device)
# # learning_rate = 1e-2
# # optimizer = torch.optim.SGD(my_module.parameters(), lr=learning_rate)
# # epoch=50
# # data =torchvision.datasets.CIFAR10(root='.\\data',train=True,transform=torchvision.transforms.Compose([torchvision.transforms.ToTensor()]),download=True)
# # train_dataloader=DataLoader(data,batch_size=64)
# # total_train_step=0
# # for i in range(epoch):
# #     my_module.train()
# #     for j in train_dataloader:
# #         img, target = j
# #         img=img.to(device)
# #         target=target.to(device)
# #         output=my_module(img)
# #         loss=loss_fn(output,target)
# #         optimizer.zero_grad()
# #         loss.backward()
# #         optimizer.step()
# #
# #         total_train_step = total_train_step + 1
# #         if total_train_step % 10 == 0:
# #             print("训练次数：{}, Loss: {}".format(total_train_step, loss.item()))
# #         if loss<2.30:
# #             break
# # torch.save(my_module,'vgg.pth')
# # print(my_module)
#
# path=r'C:\Users\Harris_Shane\Test_source\horse.jpg'
# import cv2
# pth=torch.load('vgg.pth')
# img=cv2.imread(path)
# transform=torchvision.transforms.ToTensor()
# img=transform(img)
# print(img.shape)
# # img.reshape()
# # with torch.no_grad():
# #     output = pth(img)
# print(pth)

from torch import nn
import torchvision
class vggnet(nn.Module):
    def __init__(self):
        super(vggnet,self).__init__()
        self.net=nn.Sequential(
            nn.Conv2d(3,64,3,1,1),
            nn.ReLU(),
            nn.Conv2d(64,64, 3, 1, 1),
            nn.ReLU(),
            nn.MaxPool2d(2,2),
            nn.Conv2d(64,128,3,1,1),
            nn.ReLU(),
            nn.Conv2d(128,128,3,1,1),
            nn.ReLU(),
            nn.MaxPool2d(2,2),
            nn.Conv2d(128,256,3,1,1),
            nn.ReLU(),
            nn.Conv2d(256,256,3,1,1),
            nn.ReLU(),
            nn.Conv2d(256,256,3,1,1),
            nn.ReLU(),
            nn.MaxPool2d(2,2),
            nn.Conv2d(256,512,3,1,1),
            nn.ReLU(),
            nn.Conv2d(512,512,3,1,1),
            nn.ReLU(),
            nn.Conv2d(512,512,3,1,1),
            nn.ReLU(),
            nn.MaxPool2d(2,2),
            nn.Conv2d(512,512,3,1,1),
            nn.ReLU(),
            nn.Conv2d(512,512,3,1,1),
            nn.ReLU(),
            nn.Conv2d(512,512,3,1,1),
            nn.ReLU(),
            nn.MaxPool2d(2,2),
        )
        self.classify=nn.Sequential(
            nn.Linear(512*7*7, 4096),
            nn.ReLU(),
            nn.Dropout(),
            nn.Linear(4096, 4096),
            nn.ReLU(),
            nn.Dropout(),
            nn.Linear(4096, 10),
        )
    def forward(self,x):
        x=self.net(x)
        x=x.view(x.size(0),-1)
        # print(x.size())
        x=self.classify(x)
        return x

vgg=vggnet()
import torchvision
data_=torchvision.datasets.CIFAR10("./data",train=False,transform=torchvision.transforms.Compose([torchvision.transforms.Resize(224),torchvision.transforms.ToTensor()]),download=True)
from torch.utils.data import DataLoader
data_D=DataLoader(data_,batch_size=64)
loss_fn=nn.CrossEntropyLoss()
lr=1e-2
optim=torch.optim.SGD(my_module.parameters(), lr=lr)
epoch=20
# print(len(data_D))
for i in range(epoch):
    loss_=None
    for j in data_D:
        img,lable=j
        y=vgg(img)
        loss=loss_fn(y,lable)
        optim.zero_grad()
        loss.backward()
        optim.step()
        loss_=loss
        print(loss.item())

    print('{}epoch loss is {}'.format(i,loss_))



